import numpy as np


def extract(data_transform, mode):
    """
    :param data_transform: DataTransform对象
    :param mode: accusation或relevant_articles或imprisonment
    """
    path = 'data/preprocess/'+mode
    # 创建one-hot标签
    data_transform.creat_label_set(name=mode)
    label = data_transform.creat_labels(name=mode)

    # 对于刑期还需保留连续数值
    if mode == 'imprisonment':
        path += '_dispersed'
        imprisonments = data_transform.extraction['imprisonment']
        np.save('data/preprocess/imprisonment', imprisonments)

    np.save(path, label)

